import json
import tempfile
import time
import os
import numpy as np
from collections import deque
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec
from sklearn.model_selection import train_test_split
from simulate_sensor_data import BusDevice
from extract_features import DriverScoringSystem
from main_fa2 import StateRecognizerML
import threading
import gc
import psutil


class RealTimeDriverMonitoring:
    def __init__(self):
        # 初始化配置
        self.buffer_max_size = 1500  # 缓冲池最大容量
        self.data_interval = 0.01  # 数据生成间隔(秒) 100Hz
        self.training_interval = 300  # 每300个数据点考虑训练一次模型
        self.window_size = 5  # 每次处理的数据窗口数量
        self.min_data_for_training = 5  # 训练所需的最小窗口数量
        self.model_update_counter = 0
        self.model_update_threshold = 5  # 每处理5批数据后更新模型

        # 初始化组件
        self.buffer_pool = deque(maxlen=self.buffer_max_size)
        self.scoring_system = DriverScoringSystem()
        self.recognizer = StateRecognizerML()
        self.vehicle = BusDevice("BUS001", aggressive=True)
        self.time_counter = 0
        self.running = True
        self.processing_lock = threading.Lock()
        self.last_state = "未知状态"

        # 设置matplotlib
        self._setup_matplotlib()

        # 用于存储预测结果的历史
        self.prediction_history = deque(maxlen=50)
        self.confidence_history = deque(maxlen=50)
        self.score_history = deque(maxlen=50)

        # 初始化性能计数器
        self.start_time = time.time()
        self.processed_count = 0

    def _setup_matplotlib(self):
        """配置matplotlib参数并创建监控面板"""
        plt.ion()  # 启用交互模式
        try:
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.rcParams['axes.unicode_minus'] = False
        except Exception as e:
            print(f"加载字体失败: {e}")

        # 创建监控面板
        self.fig = plt.figure(figsize=(15, 12), constrained_layout=True)
        gs = GridSpec(4, 2, figure=self.fig)

        # 加速度图表
        self.acc_ax = self.fig.add_subplot(gs[0, 0])
        self.acc_ax.set_title('三轴加速度数据')
        self.acc_ax.set_xlabel('时间步')
        self.acc_ax.set_ylabel('加速度 (g)')
        self.acc_ax.grid(True)
        self.acc_line_x, = self.acc_ax.plot([], [], label='X轴', color='#EF476F')
        self.acc_line_y, = self.acc_ax.plot([], [], label='Y轴', color='#06D6A0')
        self.acc_line_z, = self.acc_ax.plot([], [], label='Z轴', color='#118AB2')
        self.acc_ax.legend(loc='upper right')

        # 陀螺仪图表
        self.gyro_ax = self.fig.add_subplot(gs[0, 1])
        self.gyro_ax.set_title('Z轴角速度数据')
        self.gyro_ax.set_xlabel('时间步')
        self.gyro_ax.set_ylabel('角速度 (deg/s)')
        self.gyro_ax.grid(True)
        self.gyro_line_z, = self.gyro_ax.plot([], [], color='#FFD166')

        # 状态历史图表
        self.state_ax = self.fig.add_subplot(gs[1, :])
        self.state_ax.set_title('驾驶状态历史')
        self.state_ax.set_xlabel('时间窗口')
        self.state_ax.set_ylabel('状态')
        self.state_ax.grid(True)
        self.state_line, = self.state_ax.plot([], [], 'o-', markersize=8)

        # 评分历史图表
        self.score_ax = self.fig.add_subplot(gs[2, :])
        self.score_ax.set_title('驾驶评分历史')
        self.score_ax.set_xlabel('时间窗口')
        self.score_ax.set_ylabel('评分')
        self.score_ax.set_ylim(0, 100)
        self.score_ax.grid(True)
        self.score_line, = self.score_ax.plot([], [], 'o-', color='#118AB2')

        # 置信度图表
        self.conf_ax = self.fig.add_subplot(gs[3, 0])
        self.conf_ax.set_title('预测置信度')
        self.conf_ax.set_xlabel('时间窗口')
        self.conf_ax.set_ylabel('置信度')
        self.conf_ax.set_ylim(0, 1)
        self.conf_ax.grid(True)
        self.conf_line, = self.conf_ax.plot([], [], 'o-', color='#06D6A0')

        # 性能监控
        self.perf_ax = self.fig.add_subplot(gs[3, 1])
        self.perf_ax.axis('off')
        self.perf_text = self.perf_ax.text(0.05, 0.5, '', fontsize=10)

        plt.tight_layout()
        plt.pause(0.1)

    def _save_buffer_pool(self):
        """安全地保存缓冲池数据到文件"""
        target_dir = os.path.dirname(os.path.abspath('buffer_pool.json'))
        try:
            with tempfile.NamedTemporaryFile(
                    mode='w',
                    delete=False,
                    suffix='.json',
                    dir=target_dir
            ) as tmpfile:
                # 只保存最近500个数据点
                json.dump(list(self.buffer_pool)[-500:], tmpfile, indent=2)
                tmpfile.flush()
                os.fsync(tmpfile.fileno())

            # 原子替换原文件
            os.replace(tmpfile.name, 'buffer_pool.json')
            print(f"缓冲池数据已保存 (共{len(self.buffer_pool)}个数据点)")
        except Exception as e:
            print(f"保存缓冲池数据失败: {e}")

    def _update_plots(self):
        """更新所有图表"""
        # 更新加速度图表
        if len(self.recognizer.acc_x_history) > 0:
            acc_x = self.recognizer.acc_x_history[-1]
            self.acc_line_x.set_data(np.arange(len(acc_x)), acc_x)
            self.acc_line_y.set_data(np.arange(len(acc_x)), self.recognizer.acc_y_history[-1])
            self.acc_line_z.set_data(np.arange(len(acc_x)), self.recognizer.acc_z_history[-1])
            self.acc_ax.relim()
            self.acc_ax.autoscale_view()

        # 更新陀螺仪图表
        if len(self.recognizer.gyro_z_history) > 0:
            gyro_z = self.recognizer.gyro_z_history[-1]
            self.gyro_line_z.set_data(np.arange(len(gyro_z)), gyro_z)
            self.gyro_ax.relim()
            self.gyro_ax.autoscale_view()

        # 更新状态历史
        if len(self.prediction_history) > 0:
            states = list(self.prediction_history)
            state_map = {"加速": 1, "减速": 2, "转弯": 3, "平缓行驶": 4, "匀速行驶": 5, "未知状态": 0}
            state_nums = [state_map.get(s, 0) for s in states]
            self.state_line.set_data(np.arange(len(state_nums)), state_nums)
            self.state_ax.set_ylim(-0.5, 5.5)
            self.state_ax.set_yticks([0, 1, 2, 3, 4, 5])
            self.state_ax.set_yticklabels(["未知", "加速", "减速", "转弯", "平缓", "匀速"])

        # 更新评分历史
        if len(self.score_history) > 0:
            scores = list(self.score_history)
            self.score_line.set_data(np.arange(len(scores)), scores)
            self.score_ax.set_xlim(0, len(scores) + 1)

        # 更新置信度图表
        if len(self.confidence_history) > 0:
            confidences = list(self.confidence_history)
            self.conf_line.set_data(np.arange(len(confidences)), confidences)
            self.conf_ax.set_xlim(0, len(confidences) + 1)

        # 更新性能监控
        elapsed = time.time() - self.start_time
        fps = self.processed_count / elapsed if elapsed > 0 else 0
        mem_usage = psutil.Process().memory_info().rss / (1024 * 1024)  # MB
        self.perf_text.set_text(
            f"处理窗口: {self.processed_count}\n"
            f"速度: {fps:.1f} FPS\n"
            f"内存: {mem_usage:.1f} MB\n"
            f"队列: {self.recognizer.data_queue.qsize()}\n"
            f"当前状态: {self.last_state}"
        )

        # 刷新图表
        self.fig.canvas.draw_idle()
        self.fig.canvas.flush_events()

    def _process_data(self):
        """处理缓冲池中的数据并训练模型"""
        # 检查是否有足够数据
        if len(self.buffer_pool) < 100:  # 至少需要100个数据点
            return False

        # 提取传感器数据
        acc_x, acc_y, acc_z, gyro_z = self.recognizer.extract_sensor_data(list(self.buffer_pool))

        if acc_x is None or len(acc_x) == 0:
            print("提取传感器数据失败")
            return False

        # 预处理数据
        windows = self.recognizer.preprocess_data(acc_x, acc_y, acc_z, gyro_z)

        if not windows or len(windows['acceleration_x']) == 0:
            print("数据预处理后没有有效的数据窗口")
            return False

        # 准备数据集
        all_windows_data = [
            {
                'acceleration_x': windows['acceleration_x'][i],
                'acceleration_y': windows['acceleration_y'][i],
                'acceleration_z': windows['acceleration_z'][i],
                'gyroscope_z': windows['gyroscope_z'][i]
            }
            for i in range(min(len(windows['acceleration_x']), self.window_size))
        ]

        if len(all_windows_data) < self.min_data_for_training:
            print(f"数据窗口不足 ({len(all_windows_data)} < {self.min_data_for_training})，跳过训练")
            return False

        # 准备训练数据
        X, y = self.recognizer.prepare_dataset(all_windows_data)

        if len(X) <= 1:
            print("数据窗口数量不足，无法分割训练集和测试集")
            return False

        # 更新模型计数器
        self.model_update_counter += 1

        # 定期更新模型
        if self.model_update_counter >= self.model_update_threshold:
            print("\n开始更新模型...")
            self.model_update_counter = 0

            # 分割数据集并训练模型
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42
            )

            self.recognizer.train_model(X_train, y_train)
            accuracy = self.recognizer.evaluate_model(X_test, y_test)
            print(f"模型准确率: {accuracy:.2f}")

        # 启动实时处理
        self.recognizer.start_real_time_processing(all_windows_data)
        return True

    def run(self):
        """主运行循环"""
        try:
            print("启动实时驾驶员监控系统...")
            print("按 Ctrl+C 停止程序")

            last_processing_time = time.time()

            while self.running:
                start_time = time.time()

                # 1. 生成并处理传感器数据
                self.vehicle.update_sensor_data()
                data = self.vehicle.to_dict()

                # 2. 提取原始数据
                acceleration = np.array([data['accel_x'], data['accel_y'], data['accel_z']])
                gyroscope = np.array([data['gyro_x'], data['gyro_y'], data['gyro_z']])

                # 3. 计算时间间隔 - 使用固定间隔
                dt = self.data_interval

                # 4. 应用卡尔曼滤波
                self.scoring_system.kalman_filter.update(acceleration, gyroscope, dt)
                filtered_acceleration = self.scoring_system.kalman_filter.get_filtered_acceleration(acceleration)
                filtered_gyroscope = self.scoring_system.kalman_filter.get_filtered_gyroscope()

                # 5. 存储数据点
                data_point = {
                    "timestamp": self.time_counter,
                    "raw_acceleration": {"x": acceleration[0], "y": acceleration[1], "z": acceleration[2]},
                    "filtered_acceleration": {"x": filtered_acceleration[0], "y": filtered_acceleration[1],
                                              "z": filtered_acceleration[2]},
                    "raw_gyroscope": {"x": gyroscope[0], "y": gyroscope[1], "z": gyroscope[2]},
                    "filtered_gyroscope": {"x": filtered_gyroscope[0], "y": filtered_gyroscope[1],
                                           "z": filtered_gyroscope[2]}
                }
                self.buffer_pool.append(data_point)

                # 6. 定期保存数据
                if len(self.buffer_pool) % 200 == 0:  # 每200个数据点保存一次
                    self._save_buffer_pool()

                # 7. 处理数据并训练模型 - 限制处理频率
                current_time = time.time()
                if current_time - last_processing_time > 1.0:  # 每秒最多处理一次
                    if self._process_data():
                        last_processing_time = current_time

                # 8. 实时处理可视化数据
                if hasattr(self.recognizer, 'processing_running') and self.recognizer.processing_running:
                    try:
                        # 处理队列中的数据
                        self.recognizer.process_queue_data()

                        # 获取最新状态
                        if len(self.recognizer.acc_x_history) > 0:
                            latest_data = {
                                'acceleration_x': self.recognizer.acc_x_history[-1],
                                'acceleration_y': self.recognizer.acc_y_history[-1],
                                'acceleration_z': self.recognizer.acc_z_history[-1],
                                'gyroscope_z': self.recognizer.gyro_z_history[-1]
                            }

                            # 预测状态
                            state, confidence = self.recognizer.predict_states(latest_data)
                            self.last_state = state

                            # 更新历史记录
                            self.prediction_history.append(state)
                            self.confidence_history.append(confidence)

                            # 计算驾驶评分
                            #score = self.scoring_system.calculate_driver_score(state, confidence)
                            #self.score_history.append(score)

                            # 更新处理计数
                            self.processed_count += 1

                            # 更新图表
                            self._update_plots()

                    except Exception as e:
                        print(f"实时处理错误: {e}")

                # 9. 更新时间并控制频率
                self.time_counter += self.data_interval
                elapsed = time.time() - start_time
                sleep_time = max(0, self.data_interval - elapsed)
                time.sleep(sleep_time)

        except KeyboardInterrupt:
            print("\n正在停止程序...")
        finally:
            self.stop()

    def stop(self):
        """安全停止程序"""
        self.running = False
        print("正在清理资源...")

        # 停止识别器
        if hasattr(self.recognizer, 'processing_running'):
            self.recognizer.processing_running = False
            if hasattr(self.recognizer, 'thread') and self.recognizer.thread:
                self.recognizer.thread.join(timeout=1.0)

        # 保存最终数据
        self._save_buffer_pool()

        # 关闭图表
        plt.close('all')
        print("程序已安全终止")


if __name__ == "__main__":
    monitor = RealTimeDriverMonitoring()
    monitor.run()